import Image, ImageDraw
from ImageFilter import BuiltinFilter
from copy import copy

class Box(object):
  def __init__(self, x1 = 0, y1 = 0, x2 = 0, y2 = 0):
    if isinstance(x1, list): self.coords = tuple(x1)
    elif isinstance(x1, tuple): self.coords = x1
    elif isinstance(x1, Box): self.coords = copy(x1.coords)
    else: self.coords = (x1, y1, x2, y2)
  def __str__(self):
    return "Box:" + str(self.coords)
  def left(self): return self.coords[0]
  def right(self): return self.coords[2]
  def top(self): return self.coords[1]
  def bottom(self): return self.coords[3]
  def topleft(self): return (self.left(), self.top())
  def bottomright(self): return (self.right(), self.bottom())
  def width(self): return self.right() - self.left()
  def height(self): return self.bottom() - self.top()
  def size(self): return(self.width(), self.height())
  def center(self):
    x1,y1,x2,y2 = self.coords
    return ((x1 + x2) / 2, (y1 + y2) / 2)
  def offset(self, (x,y)):
    x1,y1,x2,y2 = self.coords
    self.coords = (x+x1, y+y1, x+x2, y+y2)
  def myoffset(self): return (self.left(), self.top())
  def is_inside(self, point):
    return (self.left()<=point[0]<=self.right() 
      and self.top()<=point[1]<=self.bottom())

  def expand_by(self, px, canvas_size):
    x1 = max(0, self.left() - px)
    y1 = max(0, self.top() - px)
    x2 = min(canvas_size[0] - 1, self.right() + px)
    y2 = min(canvas_size[1] - 1, self.bottom() + px)
    self.coords = (x1,y1,x2,y2)

  def collide(self, box):
    if (self.left() > box.right() or self.right() < box.left()
      or self.top() > box.bottom() or self.bottom() < box.top()):
        return False
    else: return True

  def wrap(self, box):
    x1 = min(self.left(), box.left())
    y1 = min(self.top(), box.top())
    x2 = max(self.right(), box.right())
    y2 = max(self.bottom(), box.bottom())
    self.coords = (x1, y1, x2, y2)

#
# Image Analysis Functions
#

def binarizeImage(img, threashold=64):
  return img.convert("L").point([0]*threashold+[255]*(256-threashold)).convert('1')

#def content_map(im):
  #assert im.mode == "1"
  #org_dat = im.getdata()
  #w,h = im.size
  #ret_dat = [1] * (w * h)
  #for y in range(h):
    #line_pos = y * w
    #for x in range(w):
      #pos = line_pos + x
      #if org_dat[pos] > 0: ret_dat[pos] = 0
      #else:
        #left = x
        #ret_dat[pos] = 0
        #for x in range(w-1, left, -1):
          #pos = line_pos + x
          #if org_dat[pos] != 0: ret_dat[pos] = 0
          #else: ret_dat[pos] = 0; break
        #break
  #for x in range(w):
    #for y in range(h):
      #pos = x + y*w
      #if org_dat[pos] > 0: ret_dat[pos] = 0
      #else:
        #top = y
        #ret_dat[pos] = 0
        #for y in range(h-1, top, -1):
          #pos = x + y*w
          #if org_dat[pos] != 0: ret_dat[pos] = 0
          #else: ret_dat[pos]=0; break
        #break
  #ret_img = Image.new("1", (w,h))
  #ret_img.putdata(ret_dat)
  #return ret_img

def get_shape(im):
  assert im.mode == '1'
  org_dat = im.getdata()
  w,h = im.size
  ret_dat = [0] * (w * h)
  for y in range(h):
    line_pos = y * w
    for x in range(w):
      pos = line_pos + x
      if org_dat[pos] > 0: ret_dat[pos] += 1
      else:
        left = x
        ret_dat[pos] += 1
        for x in range(w-1, left, -1):
          pos = line_pos + x
          if org_dat[pos] != 0: ret_dat[pos] += 1
          else: break
        break
  for x in range(w):
    for y in range(h):
      pos = x + y*w
      if org_dat[pos] > 0: ret_dat[pos] += 1
      else:
        top = y
        ret_dat[pos] += 1
        for y in range(h-1, top, -1):
          pos = x + y*w
          if org_dat[pos] != 0: ret_dat[pos] += 1
          else: break
        break
  for i in range(len(ret_dat)):
    if ret_dat[i]<=1: ret_dat[i]=1
    else: ret_dat[i] = 0
  ret_img = Image.new("1", (w,h))
  ret_img.putdata(ret_dat)
  return ret_img
  

def content_map(im):
  assert im.mode == "1"
  org_dat = im.getdata()
  w,h = im.size
  ret_dat = [1] * (w * h)
  for y in range(h):
    line_pos = y * w
    for x in range(w - 1):
      pos = line_pos + x
      ret_dat[pos] = 0
      if org_dat[pos] == 0 and org_dat[pos + 1]>0:
        left = x
        for x in range(w - 1, left + 1, -1):
          pos = line_pos + x
          ret_dat[pos] = 0
          if org_dat[pos] == 0 and org_dat[pos - 1] > 0: break
        break
  for x in range(w):
    for y in range(h - 1):
      pos = x + y * w
      ret_dat[pos] = 0
      if org_dat[pos] == 0 and org_dat[pos + w]>0:
        top = y
        for y in range(h - 1, top + 1, -1):
          pos = x + y * w
          ret_dat[pos] = 0
          if org_dat[pos] == 0 and org_dat[pos - w] > 0: break
        break
  ret_img = Image.new("1", (w,h))
  ret_img.putdata(ret_dat)
  return ret_img

def remove_outframe(im):
  assert im.mode == "1"
  org_dat = im.getdata()
  w,h = im.size
  ret_dat = list(org_dat)
  for y in range(h):
    line_pos = y * w
    for x in range(w - 1):
      pos = line_pos + x
      ret_dat[pos] = 1
      if org_dat[pos] == 0 and org_dat[pos + 1]>0:
        left = x
        for x in range(w - 1, left + 1, -1):
          pos = line_pos + x
          ret_dat[pos] = 1
          if org_dat[pos] == 0 and org_dat[pos - 1] > 0: break
        break
  for x in range(w):
    for y in range(h - 1):
      pos = x + y * w
      ret_dat[pos] = 1
      if org_dat[pos] == 0 and org_dat[pos + w]>0:
        top = y
        for y in range(h - 1, top + 1, -1):
          pos = x + y * w
          ret_dat[pos] = 1
          if org_dat[pos] == 0 and org_dat[pos - w] > 0: break
        break
  ret_img = Image.new("1", (w,h))
  ret_img.putdata(ret_dat)
  return ret_img
  

def get_bound_boxes(im):
  assert im.mode == "L"
  w,h = im.size
  boxes = {}
  dat = im.getdata()
  def scan_box(x, y, lbl):
    current_box = boxes.get(lbl)
    if current_box:
      x1,y1,x2,y2 = current_box.coords
      if   x < x1: current_box.coords = (x,  y1, x1, y2)
      elif x > x2: current_box.coords = (x1, y1, x,  y2)
      if   y < y1: current_box.coords = (x1, y,  x1, y2)
      elif y > y2: current_box.coords = (x1, y1, x1, y)
    else: boxes[lbl] = Box(x,y,x,y)
  for y in range(h):
    line_pos = y*w
    for x in range(w):
      pos = line_pos + x
      if dat[pos]!=0: scan_box(x, y, dat[pos]) 
  #dic -> list
  boxes = boxes.values()

  for box in boxes: box.expand_by(3, im.size)
  # box collision check
  
  for i in range(len(boxes)):
    for j in range(i+1, len(boxes)):
      if boxes[i].collide(boxes[j]):
        boxes[j].wrap(boxes[i])
        boxes[i]=None
        break
  # remove nodes or small box
  boxes = [box for box in boxes
      if box != None and box.width()>20 and box.height()>20]
  return boxes

def color_labeled(im):
  assert im.mode == "L"
  col = ((255, 0 , 0), (0, 255, 0), (0, 0, 255), 
      (0, 255, 255), (255, 0, 255), (255, 255, 0))
  ret = list(im.getdata())
  for i,val in enumerate(ret):
    if val == 0: ret[i] = (0, 0, 0)
    else: ret[i] = col[val%6]
  new_img = Image.new("RGB", im.size)
  new_img.putdata(ret)
  return new_img


def create_labeled_image(im):
  global lbls
  assert im.mode == "1"
  w,h = im.size
  draw = ImageDraw.Draw(im)
  draw.rectangle([0,0,w-1,h-1], outline=0)
  im.histogram()
  org_dat = im.getdata()
  ret_dat = [0] * (w * h)
  lbls = [0]
  neighbors = [0, 0, 0, 0]
  for y in range(1, h):
    line_pos = y * w
    for x in range(1, w-1):
      pos = line_pos + x
      if org_dat[pos]==0: continue
      neighbors[0] = ret_dat[pos - 1]
      neighbors[1] = ret_dat[pos - 1 - w]
      neighbors[2] = ret_dat[pos - w]
      neighbors[3] = ret_dat[pos + 1 - w]

      lbl = max(neighbors)
      if lbl == 0:
        new_lbl = len(lbls)
        ret_dat[pos] = new_lbl
        lbls.append(new_lbl)
      else:
        ret_dat[pos] = lbl
        for prev in neighbors:
          if prev != 0: 
            lbls[prev] = lbl
      assert ret_dat[pos]!=0, "%d, %d"%(x,y)
  #solve collisions
  final_labels = []
  for id in range(1, len(lbls)):
    if id == lbls[id]: final_labels.append(id)
    else:
      i = id
      while i != lbls[i]:
        i = lbls[i]
      lbls[id] = i
  label_dic = {0:0}
  for i,val in enumerate(final_labels): label_dic[val] = i+1
  for i in range(len(lbls)):
    lbls[i] = label_dic[lbls[i]]
  # rescan
  for i,val in enumerate(ret_dat):
    ret_dat[i] = lbls[val]
  label_im = Image.new("L", (w,h))
  label_im.putdata(ret_dat)
  return label_im

def trim_img(org_im):
  bm = content_map(binarizeImage(org_im))
  box = Box(bm.getbbox())
  box.expand_by(5, bm.size)
  return cut_img(org_im, box), box

def cut_img(im, box):
  rgn = im.crop(box.coords)
  new_img = Image.new(im.mode, (box.width(), box.height()))
  new_img.paste(rgn, (0, 0, box.width(), box.height()))
  return new_img

def separate_windows(im):
  boxes = separate_boxes(im)
  return [cut_img(im, box) for box in boxes]

def separate_boxes(original_image):
  ret = []
  def proc(bm, offset):
    assert bm.mode == "1"
    im = create_labeled_image(content_map(bm))
    boxes = get_bound_boxes(im)
    if len(boxes) == 1:
      box = boxes[0]
      box.offset(offset)
      ret.append(box)
      return
    for box in boxes:
      proc(cut_img(bm, box), (offset[0]+box.left(), offset[1]+box.top()))
  if original_image.mode != "1":
    bm = binarizeImage(original_image)
  else: bm = original_image
  proc(bm, (0,0))
  for box in ret: box.expand_by(3, original_image.size)
  return ret

def test(im):
  im1 = content_map(binarizeImage(im))
  im1_1 = im1.point(lambda x: x*255)
  im2 = create_labeled_image(im1)
  boxes = get_bound_boxes(im2)
  im3 = color_labeled(im2)
  for box in boxes:
    draw = ImageDraw.Draw(im3)
    draw.rectangle(box.coords, outline = (255,255,255))
  return [im1_1, im3, boxes]


if __name__ == '__main__':
  global im
  im = Image.open("pt.png")
  boxes = separate_boxes(im)
  for box in boxes:
    draw = ImageDraw.Draw(im)
    draw.rectangle(box.coords, outline = (255,0,0))
  im.show()
  print('finished')

  
